import akshare as ak
import pandas as pd
from sqlalchemy import create_engine
import datetime as dt
import queue
import threading
import random
import time
from bisect import bisect_right
import numpy as np
import math
import baostock as bs
from sqlalchemy import types

engine = create_engine('postgresql://quantify_qfq:xiao19911115@8.134.204.126:35432/quantify_qfq')
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)


def all():
    stock_zh_a_spot_em_df = ak.stock_zh_a_spot()
    stock_zh_a_spot_em_df['代码'] = stock_zh_a_spot_em_df['代码'].str[-6:]
    stock_zh_a_spot_em_df.to_sql('stock_spot', engine, if_exists='replace', index=False)
    return stock_zh_a_spot_em_df


def zz2000():
    # stock_cy_a_spot_em_df = ak.stock_cy_a_spot_em()
    # stock_kc_a_spot_em_df = ak.stock_kc_a_spot_em()
    zz800 = ak.index_stock_cons(symbol="000906")
    zz1000 = ak.index_stock_cons(symbol="000852")
    zz2000 = ak.index_stock_cons(symbol="932000")
    # cy_codes = set(stock_cy_a_spot_em_df['代码'])
    # kc_codes = set(stock_kc_a_spot_em_df['代码'])
    # all_exclude_codes = cy_codes | kc_codes
    df = pd.concat([zz800['品种代码'], zz1000['品种代码'], zz2000['品种代码']])
    df = df[df.str.startswith('0') | df.str.startswith('6')]
    # df = df[~df.isin(all_exclude_codes)]
    df = df.drop_duplicates()
    df.to_sql('zz2000', engine, if_exists='replace', index=False)
    return df


def producer(result_queue):
    df_all = pd.read_sql_query('select 代码,最新价,涨跌幅 from stock_spot where 今开>0 order by 代码', engine)
    zz2000 = pd.read_sql_query('select 品种代码 from zz2000', engine)
    today = dt.date.today()
    if dt.datetime.now().hour < 15:
        today -= dt.timedelta(days=1)
    if today.weekday() == 5:
        today -= dt.timedelta(days=1)
    elif today.weekday() == 6:
        today -= dt.timedelta(days=2)
    print(today)
    total = len(zz2000)
    # total = len(df_all)
    count = 0
    lg = bs.login()
    for index, row in df_all.iterrows():
        if row['代码'] in zz2000['品种代码'].values:
            count += 1
            print('爬', total, count, row.代码)
            # if np.isnan(row.总市值):
            #     continue
            table = pd.read_sql_query("select tablename from pg_tables where tablename='{}'".format(row.代码), engine)
            if len(table):
                try:
                    date = pd.read_sql_query('select max(日期) from "public"."{}"'.format(row.代码), engine)
                    if date['max'][0] < today:
                        pass
                    else:
                        continue
                except:
                    pass
            # while True:
            #     try:
            rs = bs.query_history_k_data_plus("{}.{}".format('sh' if row.代码.startswith("6") else 'sz', row.代码),
                                              "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                              start_date='1991-01-01', end_date=today.strftime('%Y-%m-%d'),
                                              frequency="d", adjustflag="2")
            data_list = []
            while (rs.error_code == '0') & rs.next():
                l = rs.get_row_data()
                for index, i in enumerate(l):
                    if i == '':
                        l[index] = 0
                data_list.append(l)
            df = pd.DataFrame(data_list,
                              columns=['日期', 'code', '开盘', '最高', '最低', '收盘', '昨收', '成交量', '成交额',
                                       'adjustflag',
                                       '换手率', 'tradestatus', '涨跌幅', 'isST'])
            df['股票代码'] = df['code'].str.split('.').str[-1]
            df['日期'] = pd.to_datetime(df['日期'])
            time.sleep(random.randint(1, 3))
            result = (df, row.代码, total, count)
            result_queue.put(result)
    result_queue.put(None)  # 结束信号
    bs.logout()


def consumer(result_queue):
    dtype_mapping = {'日期': types.DATE, '开盘': types.FLOAT, '最高': types.FLOAT, '最低': types.FLOAT,
                     '收盘': types.FLOAT, '昨收': types.FLOAT, '成交量': types.FLOAT, '成交额': types.FLOAT,
                     '换手率': types.FLOAT, '涨跌幅': types.FLOAT}
    while True:
        result = result_queue.get()
        if result is None:
            time.sleep(1)
            break
        result[0].to_sql(result[1], engine, if_exists='replace', index=False, dtype=dtype_mapping)
        print('写', result[2], result[3], result[1])
        result_queue.task_done()


def rsi(df: pd.DataFrame()):
    df['cum_sum'] = df['换手率'].cumsum()
    cum_sum_array = df['cum_sum'].values
    distances = []
    for i in range(len(df)):
        if i == 0:
            distances.append(np.nan)
        else:
            target = cum_sum_array[i] - 100
            # 使用二分查找定位目标值的位置
            j = bisect_right(cum_sum_array, target, 0, i)
            if j == 0:
                distances.append(np.nan)
            else:
                # 计算行数间隔（i - j +1 转为实际行差）
                distance = i - (j - 1)
                distances.append(distance if distance > 0 else np.nan)
    df['distance'] = distances
    df['start_idx'] = df.index - df['distance'] + 1
    df['start_idx'] = df['start_idx'].fillna(0).astype(int)
    # 确保起始索引 >=0
    df['start_idx'] = df['start_idx'].clip(lower=0)
    new = df['收盘']
    delta = new.diff(1)
    df['gain'] = delta.where(delta > 0, 0)
    df['loss'] = -delta.where(delta < 0, 0)
    df['avg_gain'] = [
        df['gain'].ewm(alpha=1 / (end - start), adjust=False).mean().iloc[end - 1]
        for start, end in zip(df['start_idx'], df.index + 1)
    ]
    df['avg_loss'] = [
        df['loss'].ewm(alpha=1 / (end - start), adjust=False).mean().iloc[end - 1]
        for start, end in zip(df['start_idx'], df.index + 1)
    ]
    rs = df['avg_gain'] / df['avg_loss']
    df['rsi'] = 100 - (100 / (1 + rs))
    return df


def mfi(df: pd.DataFrame()):
    df['cum_sum'] = df['换手率'].cumsum()
    df['下一天收盘'] = df['收盘'].shift(-1)
    cum_sum_array = df['cum_sum'].values
    distances = []
    for i in range(len(df)):
        if i == 0:
            distances.append(np.nan)
        else:
            target = cum_sum_array[i] - 100
            # 使用二分查找定位目标值的位置
            j = bisect_right(cum_sum_array, target, 0, i)
            if j == 0:
                distances.append(np.nan)
            else:
                # 计算行数间隔（i - j +1 转为实际行差）
                distance = i - (j - 1)
                distances.append(distance if distance > 0 else np.nan)

    df['distance'] = distances

    df['tp'] = (df['最高'] + df['最低'] + df['收盘']) / 3
    df['rmf'] = df['tp'] * df['成交量']
    df['tp_prev'] = df['tp'].shift(1)
    df['positive_rmf'] = df['rmf'].where(df['tp'] > df['tp_prev'], 0)
    df['negative_rmf'] = df['rmf'].where(df['tp'] < df['tp_prev'], 0)

    df['start_idx'] = df.index - df['distance'] + 1
    df['start_idx'] = df['start_idx'].fillna(0).astype(int)
    # 确保起始索引 >=0
    df['start_idx'] = df['start_idx'].clip(lower=0)
    # 使用列表推导计算每行窗口求和
    df['positive_flow'] = [
        df['positive_rmf'].iloc[start:end].sum()
        for start, end in zip(df['start_idx'], df.index + 1)
    ]
    df['negative_flow'] = [
        df['negative_rmf'].iloc[start:end].sum()
        for start, end in zip(df['start_idx'], df.index + 1)
    ]

    df['mfr'] = df['positive_flow'] / df['negative_flow']
    df['mfi'] = 100 - (100 / (1 + df['mfr']))
    return df


def value_to_quantile(s, value, interpolation='midpoint'):
    """
    输入数值反推分位点
    :param s: pd.Series 输入数据
    :param value: float 目标数值
    :param interpolation: str 插值方法（需与quantile()一致）
    :return: float 分位点p（0≤p≤1）
    """
    # 1. 去重、排序数据
    sorted_s = s.dropna().sort_values().values
    n = len(sorted_s)
    if n == 0:
        return np.nan

    # 2. 计算插入位置
    pos = np.searchsorted(sorted_s, value, side='right') - 1

    # 3. 处理边界情况（超出范围）
    if pos < 0:
        return 0.0
    elif pos >= n - 1:
        return 1.0

    # 4. 根据插值方法计算分位点
    if interpolation == 'midpoint':
        fraction = (pos + 0.5) / n  # 位置中点归一化
    elif interpolation == 'linear':
        lower = sorted_s[pos]
        upper = sorted_s[pos + 1] if pos + 1 < n else lower
        fraction = (value - lower) / (upper - lower) * (1 / (n - 1)) + pos / (n - 1)
    else:
        raise ValueError("暂不支持此插值方法")

    # 5. 转换为分位点p
    p = 1.0 * (pos + fraction) / (n - 1)  # Pandas公式 pos = 1 + (n-1)*p[4,6](@ref)
    return np.clip(p, 0.0, 1.0)


def fn(x):
    if x <= 50:
        return (100 - math.sin(x * math.pi / 100) * 50) / 100
    else:
        return math.sin(x * math.pi / 100) / 2


def history():
    data_queue = queue.Queue()

    # 创建并启动线程
    p_thread = threading.Thread(target=producer, args=(data_queue,))
    c_thread = threading.Thread(target=consumer, args=(data_queue,))

    p_thread.start()
    c_thread.start()
    p_thread.join()
    c_thread.join()


def rsi2(df, x=20):
    new = df['收盘']
    delta = new.diff(1)
    df['gain2'] = delta.where(delta > 0, 0)
    df['loss2'] = -delta.where(delta < 0, 0)
    df['avg_gain2'] = df['gain2'].ewm(alpha=1 / x, adjust=False).mean()
    df['avg_loss2'] = df['loss2'].ewm(alpha=1 / x, adjust=False).mean()
    rs = df['avg_gain2'] / df['avg_loss2']
    df['rsi2'] = 100 - (100 / (1 + rs))
    return df


def mfi2(df, x=20):
    df['tp'] = (df['最高'] + df['最低'] + df['收盘']) / 3
    df['rmf'] = df['tp'] * df['成交量']
    df['tp_prev'] = df['tp'].shift(1)
    df['positive_rmf'] = df['rmf'].where(df['tp'] > df['tp_prev'], 0)
    df['negative_rmf'] = df['rmf'].where(df['tp'] < df['tp_prev'], 0)
    positive_flow = df['positive_rmf'].rolling(x).sum()
    negative_flow = df['negative_rmf'].rolling(x).sum()
    mfr = positive_flow / negative_flow
    df['mfi2'] = 100 - (100 / (1 + mfr))
    return df


def calculate_slope(window):
    """
    计算一个窗口内数据的斜率。
    window: 包含250个数据的数组。
    """
    x = np.arange(len(window))  # 创建自变量 (0, 1, 2, ..., 249)
    # 检查窗口内有效数据的数量，避免全NaN窗口报错
    if np.isnan(window).all():
        return np.nan
    # 使用np.polyfit进行线性拟合，deg=1表示一次线性方程
    slope, _ = np.polyfit(x, window[~np.isnan(window)], 1)  # 处理可能存在的NaN值
    return slope


def chiefly():
    with open('test.txt', 'w', encoding='utf-8') as ff:
        ff.write('')
    # stock_cy_a_spot_em_df = ak.stock_cy_a_spot_em()
    # stock_kc_a_spot_em_df = ak.stock_kc_a_spot_em()
    df_all = pd.read_sql_query('select 代码,最新价,涨跌幅 from stock_spot where 今开>0 order by 代码', engine)
    zz2000 = pd.read_sql_query('select 品种代码 from zz2000', engine)
    t = 0
    f = 0
    y = 0
    k = 0
    for index, row in df_all.iterrows():
        # if row['代码'] not in stock_cy_a_spot_em_df['代码'].values and row['代码'] not in stock_kc_a_spot_em_df[
        #     '代码'].values:
        if row['代码'] in zz2000['品种代码'].values:
            try:  # where 日期>'2018-1-1'
                df = pd.read_sql_query(
                    "select 日期,开盘, 收盘, 成交额,涨跌幅,最低,最高 from \"public\".\"{}\" where 日期>'2023-1-1' order by 日期".format(
                        row.代码), engine)
            except Exception:
                continue
            df = rsi2(df)
            df['rsi2_1'] = df['rsi2'].shift()
            df['买入'] = df['收盘'].shift(-1) * (1 + 0.00525)
            # df = mfi2(df)
            # df = df[['日期', '收盘', '换手率', '涨跌幅', 'start_idx', 'rsi', 'mfi']]
            # df['r_m'] = (df['rsi2'] + df['mfi2']) / 2
            # df['r_m_1'] = df['r_m'].shift()
            # df['交易变化'] = df['换手率'] / (df['换手率'].rolling(window=5).mean())
            # df['max20'] = df['收盘'].shift(-10).rolling(window=10).max()
            print(row['代码'])
            # df['std20v'] = df['涨跌幅'].rolling(window=20).std()
            df['std20'] = df['收盘'].rolling(window=20).std()
            df['boll_down'] = df['收盘'] - 1.5 * df['std20']
            df['x'] = df['收盘'].rolling(window=250).apply(calculate_slope, raw=True)
            # df['last'] = df['收盘'].shift(-20)
            # df['last'] = df['last'].fillna(df['收盘'].iloc[-1])
            df['卖出'] = 0
            df['卖出'] = df['卖出'].astype('float64')
            df['成交额5'] = df['成交额'].rolling(window=5).mean()
            # if row['代码'] == '000089':
            #     print(df.to_string())
            #     df = df[(df['收盘'].shift() < df['boll_down']) & (df['收盘'] > df['boll_down']) & (
            #             df['换手率'] > (df['换手率'].rolling(window=5).mean()) * 1.5) & (df['r_m'] < 50) & (df['x'] > 0)]
            #     print(df.to_string())
            # (df['收盘'].shift() < df['boll_down']) & (df['收盘'] > df['boll_down']) &
            new = df[(df['成交额'] > df['成交额5']) & (df['rsi2_1'] < df['rsi2']) & (df['rsi2'] < 35) & (df['x'] > 0)]
            new = new[new['买入'].notna()]
            for index2, row2 in df.iterrows():
                if index2 in new.index:
                    for index3, row3 in df.iloc[index2 + 1:index2 + 11].iterrows():
                        if row2['收盘'] * 0.85 > row3['最低']:
                            new.loc[index2, '卖出'] = row2['收盘'] * 0.85 * (1 - 0.0063)
                            break
                        if row2['收盘'] * 1.3 < row3['最高']:
                            new.loc[index2, '卖出'] = row2['收盘'] * 1.3 * (1 - 0.0063)
                            break
                    else:
                        new.loc[index2, '卖出'] = row3['收盘']
            counts = (new['买入'] < new['卖出']).value_counts()
            t += counts.get(True, 0)
            f += counts.get(False, 0)
            y += sum(
                np.where(new['卖出'] > new['买入'],
                         (new['卖出'] - new['买入']) / new['买入'],
                         0))
            k += sum(np.where(new['卖出'] > new['收盘'], 0, (new['买入'] - new['卖出']) / new['买入']))
            with open('test.txt', 'a+', encoding='utf-8') as ff:
                ff.write('{} '.format(row['代码']))
                ff.write(new.to_string())
                ff.write('\n')
            if t + f != 0:
                print(t, t + f, t / (t + f), (y / k + 1) * t / (t + f))


def real():
    df_all = pd.read_sql_query(
        "select 代码,今开,最新价,涨跌幅,成交量,最高,最低 from stock_spot where 今开>0 order by 代码",
        engine)
    zz2000 = pd.read_sql_query('select 品种代码 from zz2000', engine)
    for index, row in df_all.iterrows():
        if row['代码'] in zz2000['品种代码'].values:
            try:
                df = pd.read_sql_query(
                    "select 日期,开盘, 收盘, 成交额,涨跌幅,最低,最高 from \"public\".\"{}\" where 日期>'2024-1-1' order by 日期".format(
                        row.代码), engine)
            except Exception:
                continue
            next_day = df['日期'].iloc[-1]
            df = rsi2(df)
            df['rsi2_1'] = df['rsi2'].shift()
            df['std20'] = df['收盘'].rolling(window=20).std()
            df['x'] = df['收盘'].rolling(window=250).apply(calculate_slope, raw=True)
            df['成交额5'] = df['成交额'].rolling(window=5).mean()
            new = df[(df['成交额'] > df['成交额5']) & (df['rsi2_1'] < df['rsi2']) & (df['rsi2'] < 35) & (df['x'] > 0)]
            if len(new) and next_day == new['日期'].iloc[-1]:
                print(row['代码'])
                print(new.to_string())
                break


if __name__ == '__main__':
    # all()
    # zz2000()
    # history()
    # chiefly()
    real()
    pass
